#a simple object-based SIR model written in Python
#Jon Zelner
#University of Michigan
#August 11, 2009

import random

class simpleSIRModel():
    
    def __init__(self, b = .1, g = .01, S = 1000, I = 1, R = 0):
        self.b = b
        self.g = g
        
        self.S = S
        self.I = I
        self.R = R
        
        self.N = S+I+R
        self.N = float(self.N)
        self.t = 0
        
        self.sList = []
        self.iList = []
        self.rList = []

        self.newIList = []

    def run(self):
        while self.I > 0:
            newI = 0
            for i in range(self.S):
                if random.random() < self.b*(self.I/self.N):
                    newI += 1
            
            recoverI = 0
            for i in range(self.I):
                if random.random() < self.g:
                    recoverI += 1
            
            self.S -= newI
            self.I += (newI - recoverI)
            self.R += recoverI
            
            self.sList.append(self.S)
            self.iList.append(self.I)
            self.rList.append(self.R)
            self.newIList.append(newI)
            print('t', self.t)
            self.t += 1
        
        return [self.sList, self.iList, self.rList, self.newIList]
